#include"StringPool.h"
#include"variable.h"
#include<vector>
#include<string>
#include<fstream>
#include<ctime>
#include"Util.h"
#include"Parameters.h"
#include"DependencyInstance.h"
#include"KBestParseForest2O.h"
#include"DependencyDecoder2O.h"
using namespace std;

DependencyDecoder2O::DependencyDecoder2O(DependencyPipe *pipe) : DependencyDecoder2OGrandChildren(pipe){}

DependencyDecoder2O::~DependencyDecoder2O(){}


vector<pair<FeatureVector*, string*> >* DependencyDecoder2O::decodeProjective(DependencyInstance *inst, std::vector<vector<vector<double> > > &probs, std::vector<vector<vector<double> > > &probs_trips, std::vector<vector<vector<double> > > &probs_sibs, std::vector<vector<vector<double> > > &probs_grands, std::vector<vector<double> > &probs_gchs, std::vector<vector<vector<vector<double> > > > &nt_probs, int K){
	vector<string*>* forms = inst->forms;

	int length = (int)(forms->size());
	vector<vector<int> >* static_types = NULL;
	if(pipe->labeled){
		static_types = getTypes(nt_probs, length);
	}

	KBestParseForest2O pf = KBestParseForest2O(0, length - 1, *inst, K);

	bool added = false;
	bool added1 = false, added2 = false;

	for(int s = 0; s < length; s++){
		for(int g = 0; g < length; g++){
			pf.add(g, s, -1, 0, 0.0);
			pf.add(g, s, -1, 1, 0.0);
		}
	}

	for(int j = 1; j < length; j++){
		for(int s = 0; s < length && s + j < length; s++){
			int t = s + j;

			double prodProb_st = probs[s][t][0];
			double prodProb_ts = probs[s][t][1];

			int type1 = pipe->labeled ? (*static_types)[s][t] : 0;
			int type2 = pipe->labeled ? (*static_types)[t][s] : 0;

			double nt_prob_s_01 = 0.0;
			double nt_prob_s_10 = 0.0;
			double nt_prob_t_00 = 0.0;
			double nt_prob_t_11 = 0.0;
			
			
			if(pipe->labeled){
				nt_prob_s_01 = nt_probs[s][type1][0][1];
				nt_prob_s_10 = nt_probs[s][type2][1][0];
				nt_prob_t_00 = nt_probs[t][type1][0][0];
				nt_prob_t_11 = nt_probs[t][type2][1][1];
			}

			vector<BestPair>* pairs = NULL;

			for(int g = 0; g < length; g++){
				if(g >= s && g <= t){
					continue;
				}
				// case when r == s
				ParseForestItem *b1 = pf.getItems(g, s, s, 0 ,1);
				ParseForestItem *c1 = pf.getItems(s, s + 1, t, 1, 1);

				double prodProb_sst = probs_trips[s][s][t] + probs_sibs[s][t][0];
				double prodProb_gst = probs_grands[g][s][t] + probs_gchs[g][t];

				pairs = pf.getKBestPairs(b1, c1);
				int length1 = (int)(pairs->size());
				for(int k = 0; k < length1; k++){
					if(((*pairs)[k]).first == -1 || ((*pairs)[k]).second == -1){
						break;
					}

					int comp1 = (*pairs)[k].first;
					int comp2 = (*pairs)[k].second;


					double bc = b1[comp1].prob + c1[comp2].prob;
					// create sibling pair
					// create parent pair: s->t and s->(start,t)
					double prob_fin = bc + prodProb_st + prodProb_sst + prodProb_gst;

					if(pipe->labeled){
						prob_fin += nt_prob_s_01 + nt_prob_t_00;
					}

					added = pf.add(g, s, s, t, type1, 0, 0, prob_fin, &(b1[comp1]), &(c1[comp2]));
					if(!added){
						break;
					}
				}
				delete(pairs);

				// case when r == t
				b1 = pf.getItems(t, s, t - 1, 0, 1);
				c1 = pf.getItems(g, t, t, 1, 1);
				
				double prodProb_tts = probs_trips[t][t][s] + probs_sibs[t][s][0];
				double prodProb_gts = probs_grands[g][t][s] + probs_gchs[g][s];

				pairs = pf.getKBestPairs(b1, c1);
				length1 = (int)(pairs->size());
				for(int k = 0; k < length1; k++){
					if(((*pairs)[k]).first == -1 || ((*pairs)[k]).second == -1){
						break;
					}

					int comp1 = (*pairs)[k].first;
					int comp2 = (*pairs)[k].second;
				
					double bc = b1[comp1].prob + c1[comp2].prob;

					// create sibling pair
					// create parent pair: s->t and s->(start,t)
					double prob_fin = bc + prodProb_ts + prodProb_tts + prodProb_gts;

					if(pipe->labeled){
						prob_fin += nt_prob_t_11 + nt_prob_s_10;
					}

					added = pf.add(g, s, t, t, type2, 1, 0, prob_fin, &(b1[comp1]), &(c1[comp2]));
					if(!added){
						break;
					}
				}
				delete(pairs);

				for(int r = s; r < t; r++){
					// First case - create sibling
					b1 = pf.getItems(g, s, r, 0 ,1);
					c1 = pf.getItems(g, r + 1, t, 1, 1);

					pairs = pf.getKBestPairs(b1, c1);
					length1 = (int)(pairs->size());

					for(int k = 0; k < length1; k++){
						if(((*pairs)[k]).first == -1 || ((*pairs)[k]).second == -1){
							break;
						}

						int comp1 = (*pairs)[k].first;
						int comp2 = (*pairs)[k].second;
				
						double bc = b1[comp1].prob + c1[comp2].prob;

						added1 = pf.add(g, s, r, t, -1, 0, 2, bc, &(b1[comp1]), &(c1[comp2]));
						added2 = pf.add(g, s, r, t, -1, 1, 2, bc, &(b1[comp1]), &(c1[comp2]));
						if(!added1 && !added2){
							break;
						}
					}
					delete(pairs);
				}

				for(int r = s + 1; r < t; r++){
					// s -> (r,t)
					b1 = pf.getItems(g, s, r, 0 ,0);
					c1 = pf.getItems(s, r, t, 0, 2);
					pairs = pf.getKBestPairs(b1, c1);
					length1 = (int)(pairs->size());

					for(int k = 0; k < length1; k++){
						if(((*pairs)[k]).first == -1 || ((*pairs)[k]).second == -1){
							break;
						}

						int comp1 = (*pairs)[k].first;
						int comp2 = (*pairs)[k].second;
				
						double bc = b1[comp1].prob + c1[comp2].prob;

						double prob_fin = bc + prodProb_st + probs_trips[s][r][t] + probs_sibs[r][t][1] + probs_grands[g][s][t] + probs_gchs[g][t];

						if(pipe->labeled){
							prob_fin += nt_prob_s_01 + nt_prob_t_00;
						}
						added = pf.add(g, s, r, t, type1, 0, 0, prob_fin, &(b1[comp1]), &(c1[comp2]));
						if(!added){
							break;
						}
					}
					delete(pairs);

					// t -> (r,s)
					b1 = pf.getItems(t, s, r, 1, 2);
					c1 = pf.getItems(g, r, t, 1, 0);

					pairs = pf.getKBestPairs(b1, c1);
					length1 = (int)(pairs->size());

					for(int k = 0; k < length1; k++){
						if(((*pairs)[k]).first == -1 || ((*pairs)[k]).second == -1){
							break;
						}

						int comp1 = (*pairs)[k].first;
						int comp2 = (*pairs)[k].second;

						double bc = b1[comp1].prob + c1[comp2].prob;

						double prob_fin = bc + prodProb_ts + probs_trips[t][r][s] + probs_sibs[r][s][1] + probs_grands[g][t][s] + probs_gchs[g][s];

						if(pipe->labeled){
							prob_fin += nt_prob_t_11 + nt_prob_s_10;
						}
						added = pf.add(g, s, r, t, type2, 1, 0, prob_fin, &(b1[comp1]), &(c1[comp2]));
						if(!added){
							break;
						}
					}
					delete(pairs);
				}

				for(int r = s; r <= t; r++){
					if(r != s){
						b1 = pf.getItems(g, s, r, 0, 0);
						c1 = pf.getItems(s, r, t, 0, 1);
						pairs = pf.getKBestPairs(b1, c1);
						length1 = (int)(pairs->size());

						for(int k = 0; k < length1; k++){
							if(((*pairs)[k]).first == -1 || ((*pairs)[k]).second == -1){
								break;
							}

							int comp1 = (*pairs)[k].first;
							int comp2 = (*pairs)[k].second;

							double bc = b1[comp1].prob + c1[comp2].prob;
							added = pf.add(g, s, r, t, -1, 0, 1, bc, &(b1[comp1]), &(c1[comp2]));
							if(!added){
								break;
							}
						}
						delete(pairs);
					}

					if(r != t){
						b1 = pf.getItems(t, s, r, 1, 1);
						c1 = pf.getItems(g, r, t, 1, 0);
						pairs = pf.getKBestPairs(b1, c1);
						length1 = (int)(pairs->size());
						for(int k = 0; k < length1; k++){
							if(((*pairs)[k]).first == -1 || ((*pairs)[k]).second == -1){
								break;
							}

							int comp1 = (*pairs)[k].first;
							int comp2 = (*pairs)[k].second;

							double bc = b1[comp1].prob + c1[comp2].prob;
							added = pf.add(g, s, r, t, -1, 1, 1, bc, &(b1[comp1]), &(c1[comp2]));
							if(!added){
								break;
							}
						}
						delete(pairs);
					}
				}
			}
		}
	}

	if(static_types != NULL){
		delete(static_types);
	}
	bool normal = (pipe->options->root_type == 0);
	vector<pair<FeatureVector*, string*> >* result = pf.getBestParses(inst, pipe, probs, probs_trips, probs_sibs, nt_probs, pipe->labeled, normal);
	return result;
}